#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""

@Time   :2025/8/6 上午10:14
@Author :zengjiahao1989@gmail.com
@File   :2.JsonOutputParser.py
"""
import os

import dotenv
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

dotenv.load_dotenv()


# 1. 创建一个json数据结构，用于告诉大语言模型这个json的结构
class Joke(BaseModel):
    joke: str = Field(description='回答用户的冷笑话')

    punchline: str = Field(description='这个冷笑话的笑点')


parser = JsonOutputParser(pydantic_object=Joke)

# 2.构建一个提示模版
prompt = ChatPromptTemplate.from_template("请根据用户的提问进行回答，\n{format_instructions}\n{query}").partial(
    format_instructions=parser.get_format_instructions())

# print(prompt.format(query="请讲一个关于程序员的冷笑话"))

# 3.构建大语言模型
llm = ChatOpenAI(
    model_name="kimi-k2-0711-preview",
    openai_api_key=os.getenv("OPENAI_API_KEY"),
    openai_api_base=os.getenv("OPENAI_API_BASE"),
)

# 4.传递提示并进行解析
joke = parser.invoke(llm.invoke(prompt.invoke({"query": "请讲一个关于程序员的冷笑话"})))

print(type(joke))
print(joke.get("punchline"))
print(joke)
